import numpy as np
import pandas as pd

import networkx as nx
import matplotlib.pyplot as plt

from Config import Config
from data.data_parser import GraphParser
from data.data_set import gpu_setup

"""
画图展示mRNA分子的样子
"""


def plat_mRNA(Graph, sample):
    # 对mRNA分子进行画图
    plt.figure(figsize=(24, 12))
    nx_G = nx.Graph()
    nx_G.add_nodes_from([i for i in range(len(Graph['nfeat']))])

    nx_G.add_edges_from(Graph['edges'])
    # 颜色设置：碱基节点用绿色，密码子节点用黄色
    node_color = ['g' for _ in range(sample['seq_length'])] + \
                 ['y' for _ in range(len(Graph['nfeat']) - sample['seq_length'])]
    options = {
        "node_color": node_color,
    }
    pos = nx.spring_layout(nx_G, iterations=400, k=0.2)
    nx.draw(nx_G, pos, **options)
    plt.show()


def get_gdata(data_file, pos=0):
    """选择文件，需要显示数据的位置"""
    # 以下是基本参数配置，画图时使用不到
    model_name = 'GCN'
    device = gpu_setup(True, 0)
    config = Config(model_name, device)
    parser = GraphParser(config, mode='train')

    df = pd.read_json(data_file, lines=True)
    sample = df.loc[pos]
    # 向graph_list中添加图
    gdata = parser.parse(sample)

    return gdata, sample


if __name__ == '__main__':
    # 文件位置
    data_file = "../dataset/valid.json"
    # 要显示的数据下标
    pos = 0
    gdata, sample = get_gdata(data_file, pos)
    print(gdata)
    plat_mRNA(gdata, sample)
